const fs = require('fs');
const labelDef = require('./Def');
const BaseLabelLoader = require('../../BaseLabelLoader');
const { getImageSize } = require('../../../utils/imageUtils');

class YOLOLabelLoader extends BaseLabelLoader {
    async load(label_info) {
        // console.log('label_info = ', label_info);
        try {
            const imagePath = label_info.image_path;
            const { width: imageWidth, height: imageHeight } = getImageSize(imagePath);
            // TODO: 请补充获取图片宽高的方法
            const txtData = fs.readFileSync(label_info.label_path, 'utf8');
            const lines = txtData.trim().split('\n');
            let result = [];
            for (let i = 0; i < lines.length; i++) {
                const parts = lines[i].split(' ').map((x) => parseFloat(x));
                if (parts.length !== 5) {
                    console.warn(`Invalid line format in file: ${label_info.label_path}, line: ${i}`);
                    continue;
                }
                const [class_id, center_x, center_y, width, height] = parts;
                // console.log(`lines = `, parts);

                // 加载yolo数据集的categories
                const categories = this.app.db && (await this.app.db.categories.getCategories(label_info.dataset_id));
                const category = categories && categories.find((el) => el.category_id === class_id);

                // 构建标注数据对象
                let jsonObj = labelDef();
                jsonObj.label_format = 'yolo';
                jsonObj.id = i;
                jsonObj.category_id = class_id;
                jsonObj.category = category ? category.category_name : null;

                // Convert normalized values to actual image coordinates (if image dimensions known)
                jsonObj.bbox = [
                    Math.round((center_x - width / 2) * imageWidth),
                    Math.round((center_y - height / 2) * imageHeight),
                    Math.round((center_x + width / 2) * imageWidth),
                    Math.round((center_y + height / 2) * imageHeight)
                ];
                jsonObj.area = width * imageWidth * height * imageHeight;
                // console.log('jsonObj = ', jsonObj);
                result.push(jsonObj);
            }
            return result;
        } catch (error) {
            console.error('Error load label ', label_info, '\nerror: ', error);
            return { code: -1, msg: error.message };
        }
    }

    // 保存数据
    async save(label_info) {
        // console.log('YOLOLabelLoader::save label_info = ', label_info);
        let yoloPath = label_info.label_url;
        let yoloText = await this.fromStudio(label_info);
        fs.writeFileSync(yoloPath, yoloText.join('\n'));

        // 更新Datasets和ListDB的信息
        let dataset_id = label_info.dataset_id;
        this.app.db && this.app.db.datasets.updateDataset({ id: dataset_id });
        let list_db = this.app.list_dbs[dataset_id];
        list_db && list_db.pairs.updateLabelPath(label_info.pairs_id);
    }

    async toString(label_info, categories = null) {
        let yoloJson = await this.fromStudio(label_info, categories);
        return yoloJson.join('\n');
    }

    // 转换数据
    async fromStudio(label_info, categories = null) {
        // console.log('YOLOLabelLoader::fromStudio label_info = ', label_info);
        let dataset_id = label_info.dataset_id;
        let labels = label_info.labels;
        let image_url = label_info.image_url;

        // 读取图片的宽高，更新到voc的xml里面
        const { width: imageWidth, height: imageHeight } = getImageSize(image_url);

        // 获取categories信息
        if (!categories) {
            categories = await this.app.db.categories.getCategories(dataset_id);
        }

        // 保存yolo格式数据
        let yoloLines = [];
        for (let i = 0; i < labels.length; i++) {
            if (!labels[i]) continue;
            let jsonObj = labels[i].data;
            let yoloLine = [];

            // 处理category;
            let category = categories.find((el) => el.category_name === jsonObj.category);
            if (category) {
                jsonObj.category_id = category.category_id;
            }
            yoloLine[0] = jsonObj.category_id;

            // 处理bbox
            let xmin = jsonObj.bbox[0];
            let ymin = jsonObj.bbox[1];
            let xmax = jsonObj.bbox[2];
            let ymax = jsonObj.bbox[3];
            yoloLine[1] = (xmin + xmax) / (2 * imageWidth);
            yoloLine[2] = (ymin + ymax) / (2 * imageHeight);
            yoloLine[3] = (xmax - xmin) / imageWidth;
            yoloLine[4] = (ymax - ymin) / imageHeight;
            // console.log('jsonObj.bbox = ', jsonObj.bbox);
            // console.log('yoloLine = ', yoloLine);

            // 插入新增对象
            yoloLines.push(yoloLine.join(' '));
        }
        return yoloLines;
    }
}

module.exports = YOLOLabelLoader;
